#include "qrender/q3ds_model.h"

#include <fstream>

#include <QOpenGLContext>
#include <QOpenGLFunctions>

#include <QDebug>
#include <QDir>
#include <QtMath>
#include <QDomDocument>

#include "q3ds_process.h"

Q3dsModel::Q3dsModel()
{
    // demo 默认 DH 参数
    std::unordered_map<std::string, std::vector<double>> dh;
    dh["a"] = std::vector<double>{ 0, 0, 0.408, 0.376, 0, 0 };
    dh["alpha"] =
        std::vector<double>{ 0, -M_PI / 2., M_PI, M_PI, -M_PI / 2., M_PI / 2. };
    dh["d"] = std::vector<double>{ 0.122, 0.1215, 0, 0, 0.1025, 0.094 };
    dh["theta"] = std::vector<double>{ M_PI, -M_PI / 2., 0, -M_PI / 2., 0, 0 };

    reload(QStringList() << ".", "demo", dh);
}

bool Q3dsModel::reload(
    const QStringList &dirs, const QString &robot_type,
    const std::unordered_map<std::string, std::vector<double>> &dh)
{
    if (dirs_ == dirs && robot_type_ == robot_type && dh_ == dh) {
        return true;
    }
    dirs_ = dirs;
    robot_type_ = robot_type;
    dh_ = dh;

    // Create and load the link model
    QString urdf_dir, dir_prefix;
    for (auto &_dir : dirs) {
        auto path =
            _dir +
            QString("/3ds_model_description/urdf/%1.urdf").arg(robot_type);
        if (QFile::exists(path)) {
            urdf_dir = path;
            dir_prefix = _dir;
            break;
        }
    }

    if (urdf_dir.isEmpty()) {
        qWarning() << robot_type << ".urdf not found.";
        return false;
    }

    auto add_link = [this, &dir_prefix](const QString &link, double a,
                                        double alpha, double d, double theta) {
        links_.append(std::make_shared<Q3dsProcess>(
            QString("%1/%2").arg(dir_prefix, link.mid(10, link.size())), a,
            alpha, d, theta));
    };
    auto a = dh.at("a");
    auto alpha = dh.at("alpha");
    auto d = dh.at("d");
    auto theta = dh.at("theta");

    links_.clear();
    meshs_.clear();

    QFile file(urdf_dir);
    QDomDocument xml_doc;
    if (file.open(QIODevice::ReadOnly)) {
        if (xml_doc.setContent(&file)) {
            auto root = xml_doc.documentElement();
            auto links = root.elementsByTagName("link");
            for (int i = 0; i < links.count(); i++) {
                auto node = links.at(i);

                auto visual = getElementByName(node, "visual");
                if (!visual.isNull()) {
                    auto geometry = getElementByName(visual, "geometry");
                    if (!geometry.isNull()) {
                        auto mesh = getElementByName(visual, "mesh");
                        if (!mesh.isNull()) {
                            auto path = mesh.attribute("filename");
                            if (!path.endsWith(".3ds")) {
                                auto index = path.lastIndexOf(".");
                                path.replace(index, path.size() - index,
                                             ".3ds");
                            }
                            if (path.contains("link0")) {
                                add_link(path, 0., 0, 0, 0);
                            } else {
                                add_link(path, a[i - 1] * 1000.,
                                         qRadiansToDegrees(alpha[i - 1]),
                                         d[i - 1] * 1000.,
                                         qRadiansToDegrees(theta[i - 1]));
                            }
                        }
                    }
                }
            }
        } else {
            file.close();
            qWarning() << "Failed to load " << robot_type << ".urdf";
            return false;
        }

        file.close();
    } else {
        qWarning() << "Failed to open " << robot_type << ".urdf";
        return false;
    }

    for (auto &l : links_) {
        l->loadModel();
    }

    meshs_changed_ = true;

    return true;
}

void Q3dsModel::setJointPositions(const std::vector<double> &joint_positions)
{
    angles_.resize(joint_positions.size());
    for (size_t i = 0; i < angles_.size(); i++) {
        angles_[i] = qRadiansToDegrees(joint_positions[i]);
    }

    angles_chagned_ = true;
}

QVector<Q3dsMesh *> Q3dsModel::getMeshs()
{
    if (angles_.size() != (links_.size() - 1)) {
        return {};
    }

    // 如果数据获取过，则不再重复获取
    // 关节数据变化时，setJointPoisitions 中会立即变化数据
    if (meshs_.size() == links_.size() && !angles_chagned_) {
        return meshs_;
    }
    meshs_.clear();

    // 模型偏移
    QMatrix4x4 translation_matrix;
    translation_matrix.rotate(-90.0, 1.0, 0.0, 0.0); // Y 轴向上
    translation_matrix.rotate(-90.0, 0.0, 0.0, 1.0);

    // 底座
    auto mesh = links_[0]->getMesh(0, translation_matrix);
    meshs_ << mesh;

    // 关节
    for (size_t i = 1; i <= angles_.size(); i++) {
        mesh = links_[i]->getMesh(angles_[i - 1], translation_matrix);
        meshs_ << mesh;
    }

    angles_chagned_ = false;
    meshs_changed_ = false;

    return meshs_;
}

bool Q3dsModel::meshsChanged()
{
    return meshs_changed_;
}

QDomElement Q3dsModel::getElementByName(QDomNode node, const QString &name)
{
    if (node.isNull()) {
        return QDomElement();
    }

    auto e = node.toElement();
    if (e.isNull()) {
        return QDomElement();
    }

    auto nodes = e.elementsByTagName(name);
    if (nodes.isEmpty()) {
        return QDomElement();
    }

    return nodes.at(0).toElement();
}
